import cv2
import os
import sys
import numpy as np
from tqdm import tqdm

STRIDE_SIZE = 20 # 步长

# Image pre-process
# src_img_path：源图片存放路径，请在该路径下再建立存放源图片的文件夹
# dst_img_path：处理后的图片存放路径
# bgr_mean_file：保存图片bgr均值的文件名
def preprocess_imgs(src_img_path, dst_img_path, bgr_mean_file):
    f = open(bgr_mean_file, 'w')
    img_count = 0
    forders = os.listdir(src_img_path)
    resize_w, resize_h = STRIDE_SIZE, STRIDE_SIZE
    for forder in forders:
        imgs = os.listdir(os.path.join(src_img_path, forder))
        for i in tqdm(imgs):
            img = cv2.imread(os.path.join(os.path.join(src_img_path, forder), i))
            try:
                h, w, _ = img.shape
            except:
                continue
            if w / h > 1.5 or h / w > 1.5:
                continue
            img = cv2.resize(img, (resize_w, resize_h))
            bgr_mean = img.sum(axis=1).sum(axis=0) / (resize_h * resize_w)
            cv2.imwrite(os.path.join(dst_img_path, str(img_count).zfill(7)+'.jpg'), img)
            f.write(str(img_count).zfill(7)+'.jpg' + ' ' + ' '.join([str(int(val)) for val in bgr_mean]) + '\n')
            img_count += 1
    f.close()

# Draw a image
# dst_img_path：预处理后的图片存放路径
# src_img：要拼凑的目标源图
# bgr_mean_file：保存图片bgr均值的文件名
def draw_mask(dst_img_path, src_img, bgr_mean_file):
    stride = STRIDE_SIZE
 
    f = open(bgr_mean_file, 'r')
    labels = f.readlines()
    labels_dict = dict()
    for label in labels:
        labels_dict[label.split(' ')[0]] = [float(val) for val in label.strip().split(' ')[1:]]
    # print(labels_dict)

    img = cv2.imread(src_img)
    h, w, c = img.shape
    dst_img = np.zeros((h, w, c), dtype=np.uint8)
    for top in tqdm(range(0, h, stride)):
        if h - top < stride:
            continue
        for left in range(0, w, stride):
            if w - left < stride:
                continue
            crop_img = img[top:top+stride, left:left+stride, :]
            bgr_mean = crop_img.sum(axis=1).sum(axis=0) / (stride ** 2)
 
            min_dis = sys.maxsize
            hit_key = str()
            for key,val in labels_dict.items():
                dis = ((bgr_mean - val) ** 2).sum()
                if min_dis > dis:
                    min_dis = dis
                    hit_key = key
            # print("labels_dict len:", len(labels_dict), "min_dis:", min_dis, "dis:", dis, "hit_key:", hit_key)
            hit_img = cv2.imread(os.path.join(dst_img_path, hit_key))
            # labels_dict.pop(hit_key) # If you wanna every picture is unique, please uncomment this line to remove the hit picture 
 
            dst_img[top:top+stride, left:left+stride, :] = hit_img
 
    cv2.imwrite('dst_face.png', dst_img) # 保存为dst_face.png

preprocess_imgs("src_images", "rectified_images", "bgr_mean_file.txt")
draw_mask("rectified_images", "original_face.jpg", "bgr_mean_file.txt")